//
//  Problem1314.swift
//  TestProject
//
//  Created by 毕武侠 on 2021/2/21.
//  Copyright © 2021 zhulong. All rights reserved.
//

import UIKit

/*
 1314. 矩阵区域和
 给你一个 m * n 的矩阵 mat 和一个整数 K ，请你返回一个矩阵 answer ，其中每个 answer[i][j] 是所有满足下述条件的元素 mat[r][c] 的和：

 i - K <= r <= i + K, j - K <= c <= j + K
 (r, c) 在矩阵内。
  
 示例 1：
     输入：mat = [[1,2,3],[4,5,6],[7,8,9]], K = 1
     输出：[[12,21,16],[27,45,33],[24,39,28]]
 示例 2：
     输入：mat = [[1,2,3],[4,5,6],[7,8,9]], K = 2
     输出：[[45,45,45],[45,45,45],[45,45,45]]
  
 提示：
     m == mat.length
     n == mat[i].length
     1 <= m, n, K <= 100
     1 <= mat[i][j] <= 100
 */
@objcMembers class Problem1314: NSObject {
    func solution() {
        print(matrixBlockSum([[1,2,3],[4,5,6],[7,8,9]], 1))
        print(matrixBlockSum([[1,2,3],[4,5,6],[7,8,9]], 2))
    }
    
    /*
     数组前缀 + 动态规划
     1: 创建一个数组dp[m][n]
     2: dp[i][j]: 是0～i, 0~j的和
     3: dp[i][j] = dp[i-1][j] + dp[i][j-1] - dp[i-1][j-1](多加了一遍) + mat[i][j]
     4: 创建一个数组result[len][len]
     5: result[r][c]: i - K <= r <= i + K, j - K <= c <= j + K
        result[r][c] = dp[i+k][j+k] - dp[i+k][j-k] - dp[j-k][j+k] + dp[i-k][j-k]
     */
    func matrixBlockSum(_ mat: [[Int]], _ K: Int) -> [[Int]] {
        let m = mat.count
        let n = mat[0].count
        var dp = Array(repeating: Array(repeating: 0, count: n+1), count: m+1)
        // 计算0～m, 0~n的前缀和
        for i in 1...m {
            for j in 1...n {
                dp[i][j] = dp[i-1][j] + dp[i][j-1] - dp[i-1][j-1] + mat[i-1][j-1]
            }
        }
        
        var iMin = 0
        var iMax = 0
        var jMin = 0
        var jMax = 0
        var result = Array(repeating: Array(repeating: 0, count: n), count: m)
        for i in 0..<m {
            for j in 0..<n {
                iMin = (i - K ) <= 0 ? 0 : (i - K)
                iMax = (i + K + 1) >= m ? m : (i + K + 1)
                jMin = (j - K ) <= 0 ? 0 : (j - K)
                jMax = (j + K + 1) >= n ? n : (j + K + 1)
                result[i][j] = dp[iMax][jMax] - dp[iMax][jMin] - dp[iMin][jMax] + dp[iMin][jMin]
            }
        }
        
        return result
    }
}
